// 3. 被观察者（密钥管理）
use crate::key_manager::algorithm::create_algorithm;
use crate::key_manager::cache::entity::key_pair::KeyPair;
use crate::key_manager::cache::store::KeyStore;
use crate::key_manager::error::KeyManagerError;
use crate::key_manager::lifecycle::key_observer::observer_init::register::OBSERVER_REGISTRY;
use crate::key_manager::lifecycle::key_observer::KeyLifecycleObserver;
use crate::key_manager::model::KeyVersionModel;
use crate::key_manager::model::{VaultResponse, Version};
use crate::key_manager::orm::entity::key_manager_key_version::get_current_key_version;
use crate::key_manager::orm::entity::key_manager_key_version::update_key_version;
use crate::key_manager::util::encode_utils::decode_from_string;
use attestation_common::distributed_lock::Lock as DistributedLock;
use attestation_common::log::{error, info};
use attestation_common::rdb::get_connection;
use once_cell::sync::OnceCell;
use sea_orm::TransactionTrait;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use mockall::automock;

#[derive(Debug)]
pub struct KeySubject {
    observers: Vec<Arc<dyn KeyLifecycleObserver + Send + Sync>>,
}

impl KeySubject {
    pub fn new() -> Self {
        KeySubject {
            observers: Vec::new(),
        }
    }

    // Add watchers
    pub fn attach(&mut self, observer: Arc<dyn KeyLifecycleObserver + Send + Sync>) {
        self.observers.push(observer);
    }

    // Inform all observers
    async fn notify_observers(
        &self,
        key_version: &str,
    ) -> Result<(), KeyManagerError> {
        let registry = OBSERVER_REGISTRY.get().unwrap();
        let observers = registry.lock().unwrap();
        let db_connection = get_connection().await.unwrap();
        // Open a transaction
        let tx = db_connection.begin().await.unwrap();
        let tx_arc = Arc::new(tx);
        let mut tx_holder = Some(tx_arc);
        for observer in observers.iter() {
            info!("KeySubject: Notifying observer: ");
            let current_tx = Arc::clone(tx_holder.as_ref().unwrap());
            match observer.signature_update(key_version, current_tx).await {
                Ok(_) => (),
                Err(_e) => {
                    error!("KeySubject: Failed to notify observers");
                    // Roll back the transaction
                    if let Some(tx_arc) =  tx_holder.take() {
                        match Arc::try_unwrap(tx_arc) {
                            Ok(_tx) => (),
                            Err(_e) => {
                                return Err(KeyManagerError::new("Unable to rollback the transaction: there are unreleased references.".to_string()));
                            }
                        }
                    }
                    return Err(KeyManagerError::new(
                        "Failed to notify observers".to_string(),
                    ));
                },
            }
        }
        // 提交事务时需要从 Arc 中取出原始对象
        if let Some(tx_arc) = tx_holder.take() {
            match Arc::try_unwrap(tx_arc) {
                Ok(tx) => tx.commit().await?,
                Err(_) => return Err(KeyManagerError::new("Unable to submit the transaction: there are unreleased references.".to_string())),
            }
        }
        info!("KeySubject: Notified observers");
        Ok(())
    }
}

#[automock]
pub trait KeyLifecycle {
    /**
     * 密钥轮换
     *
     * @param vault_response 全量密钥
     * @return 初始化结果
     */
    fn perform_key_rotation(
        &self,
        vault_response: VaultResponse,
    ) -> impl std::future::Future<Output = Result<(), KeyManagerError>>;

    /**
     * 生成密钥对
     *
     * @param vault_response 全量密钥
     * @param db_version 数据库密钥版本
     * @return 初始化结果
     */
    fn generate_key_pair(
        vault_response: &VaultResponse,
        db_version: &Version,
    ) -> impl std::future::Future<Output = Result<(), KeyManagerError>>;
}

impl KeyLifecycle for KeySubject {
    fn perform_key_rotation(
        &self,
        vault_response: VaultResponse,
    ) -> impl std::future::Future<Output = Result<(), KeyManagerError>> {
        async move {
            info!("KeySubject: Performing key rotation...");
            // Check the database key version
            let db_version = get_current_key_version().await.unwrap_or_else(|e| {
                info!("No current version found, using v1 as default: {}", e);
                Version::new("v1")
            });
            info!("Database key version: {}", db_version);
            // Derived key database version v1: Query the latest key v5: Derived key range i v1-v5
            Self::generate_key_pair(&vault_response, &db_version).await?;

            let max_version = vault_response.find_max_version().unwrap();
            info!("Max version: {}", max_version);
            if db_version >= max_version {
                info!(
                    "db_version is large than max_version, db_version: {}, max_version: {}",
                    db_version, max_version
                );
                return Ok(());
            }
            let key_store = KeyStore::global();
            let latest_version = key_store.get_latest_version().unwrap();
            info!("KeySubject: latest_version: {}", latest_version);
            // Get distributed locks
            KeyRotationLock::new().acquire()?;
            info!("KeySubject: Acquired key rotation lock");
            // Check the database version again
            let db_version = get_current_key_version().await.unwrap_or_else(|e| {
                info!("No current version found, using v1 as default: {}", e);
                Version::new("v1")
            });
            info!("KeySubject: Database key version: {}", db_version);
            // Whether rotation is required
            let need_rotation = db_version < max_version;
            if need_rotation {
                // key rotation
                info!("KeySubject: Key rotation is needed. Notifying observers...");
                let max_version = max_version.to_string();
                let notify_result = self.notify_observers(max_version.as_str());
                if notify_result.await.is_err() {
                    error!("KeySubject: Failed to notify observers");
                    // Roll back the transaction
                    return Err(KeyManagerError::new(
                        "Failed to notify observers".to_string(),
                    ));
                }
                info!("max_version.to_string(): {}", max_version.to_string());
                // Update the database key version
                // Open a transaction
                let db_connection = get_connection().await.unwrap();
                let tx = db_connection.begin().await.unwrap();
                update_key_version(max_version.to_string(), &tx).await?;
                tx.commit().await?;
            }
            Ok(())
        }
    }

    fn generate_key_pair(
        vault_response: &VaultResponse,
        db_version: &Version,
    ) -> impl std::future::Future<Output = Result<(), KeyManagerError>> {
        async move {
            info!("KeySubject: Generating key pair...");
            let store = KeyStore::global();
            let process =
                |(key_type, keys): (&str, &Vec<KeyVersionModel>)| -> Result<(), KeyManagerError> {
                    let keys = keys.iter().filter(|v| &v.version >= db_version);

                    keys.for_each(|v| {
                        init_store(&store, key_type, v).unwrap();
                        info!("KeySubject: Initialized store for key type: {}", key_type);
                    });
                    Ok(())
                };
            [
                ("FSK", &vault_response.fsk),
                ("NSK", &vault_response.nsk),
                ("ASK", &vault_response.ask),
            ]
            .into_iter()
            .try_for_each(process)
        }
    }
}

fn init_store(
    store: &KeyStore,
    key_type: &str,
    v: &KeyVersionModel,
) -> Result<(), KeyManagerError> {
    info!("KeySubject: Initializing store for key type: {}", key_type);
    let private_key = decode_from_string(&v.private_key, &v.encoding).unwrap();
    let (private, public) = create_algorithm(&v.algorithm).and_then(|alg| {
        alg.derive_public(&private_key)
            .map(|pub_key| (private_key.clone(), pub_key))
            .or_else(|_| {
                panic!("Failed to derive public key for key type: {}", key_type);
            })
    })?;
    store.insert(
        key_type,
        &v.version,
        KeyPair {
            cached_private: OnceCell::new(),
            cached_public: OnceCell::new(),
            private_bytes: private,
            public_bytes: public,
            algorithm: v.algorithm.clone().to_string(),
        },
    )
}

struct KeyRotationLock {
    lock: DistributedLock,
}

impl KeyRotationLock {
    pub fn new() -> Self {
        Self {
            lock: DistributedLock::new("key_rotation_lock", "key_rotation_lock", 1, 1).unwrap(),
        }
    }

    pub fn acquire(&self) -> Result<(), KeyManagerError> {
        loop {
            match self.lock.acquire() {
                Ok(true) => {
                    info!("KeySubject: Acquired key rotation lock");
                    return Ok(());
                }
                Ok(false) => {
                    info!("Failed to acquire key rotation lock, retrying...");
                    thread::sleep(Duration::from_secs(5));
                    continue;
                }
                Err(e) => {
                    error!("Failed to acquire key rotation lock: {}", e);
                    return Err(KeyManagerError::new(e.to_string()));
                }
            }
        }
    }
}

impl Drop for KeyRotationLock {
    fn drop(&mut self) {
        if let Err(e) = self.lock.release() {
            error!("Failed to release key rotation lock: {}", e);
        }
    }
}

#[cfg(test)]
#[allow(warnings)]
mod tests {
    use super::*;
    use crate::key_manager::error::KeyManagerError;
    use crate::key_manager::model::{Algorithm, Encoding, PrivateKey};
    use mockall::mock;
    use mockall::predicate::*;
    use sea_orm::{DatabaseBackend, DatabaseTransaction, MockDatabase, MockExecResult};
    use std::collections::HashMap;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::{Arc, RwLock};
    use crate::key_manager::lifecycle::key_observer::key_lifecycle_observer::BoxFuture;

    // ---------- Mock Redis Client 实现 ----------
    mock! {
        pub MockLockClient {
            fn acquire_lock(&self, key: &str, value: &str, ttl: u64) -> Result<bool, String>;
            fn release_lock(&self, key: &str, value: &str) -> Result<bool, String>;
        }
    }

    type VersionMap = Arc<RwLock<HashMap<String, KeyPair>>>;

    // ---------- Mock KeyStore 配置 ----------
    struct TestKeyStore {
        fsk: VersionMap,
        nsk: VersionMap,
        ask: VersionMap,
    }

    impl TestKeyStore {
        fn new() -> Self {
            Self {
                fsk: Arc::new(RwLock::new(HashMap::new())),
                nsk: Arc::new(RwLock::new(HashMap::new())),
                ask: Arc::new(RwLock::new(HashMap::new())),
            }
        }

        fn to_keystore(&self) -> KeyStore {
            KeyStore {
                inner: HashMap::from([
                    ("FSK".to_string(), self.fsk.clone()),
                    ("NSK".to_string(), self.nsk.clone()),
                    ("ASK".to_string(), self.ask.clone()),
                ]),
                latest_version: OnceCell::new(),
            }
        }
    }

    // ---------装的观察者实现 ----------
    #[derive(Debug)]
    struct MockObserver {
        called: AtomicBool,
    }

    impl MockObserver {
        fn new() -> Self {
            Self {
                called: AtomicBool::new(false),
            }
        }
    }

    impl KeyLifecycleObserver for MockObserver {
        fn signature_update(&self, key_version: &str, tx: Arc<DatabaseTransaction>) -> BoxFuture<Result<(), Box<KeyManagerError>>> {
            Box::pin(async move {Ok(())})
        }
    }

    // ---------- 测试用例 ----------
    #[tokio::test]
    async fn test_key_rotation_success() {
        let _test_keystore = TestKeyStore::new();
        let db = MockDatabase::new(DatabaseBackend::Postgres).into_connection();
        let _tx = db.begin().await.unwrap();

        let vault_response = VaultResponse {
            fsk: vec![],
            nsk: vec![],
            ask: vec![KeyVersionModel {
                version: Version::new("v2"),
                private_key: PrivateKey("priv_key".into()),
                algorithm: Algorithm("RS256".into()),
                encoding: Encoding("base64".into()),
            }],
        };

        let mut mock_lock = MockMockLockClient::new();
        mock_lock
            .expect_acquire_lock()
            .returning(|_, _, _| Ok(true));

        let subject = KeySubject::new();
        let result = subject.perform_key_rotation(vault_response).await;

        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_no_rotation_when_version_up2date() {
        let test_db = MockDatabase::new(DatabaseBackend::Postgres);
        test_db.append_exec_results(vec![MockExecResult::default()]);

        let vault_response = VaultResponse {
            fsk: vec![KeyVersionModel {
                version: Version::new("v2"),
                private_key: PrivateKey("priv_key".into()),
                algorithm: Algorithm("RS256".into()),
                encoding: Encoding("base64".into()),
            }],
            nsk: vec![],
            ask: vec![],
        };

        let observer = Arc::new(MockObserver::new());
        let mut subject = KeySubject::new();
        subject.attach(observer.clone());

        let result = subject.perform_key_rotation(vault_response).await;

        assert!(result.is_ok());
        assert!(!observer.called.load(Ordering::SeqCst));
    }

    #[test]
    fn test_duplicate_key_insert() {
        let keystore = TestKeyStore::new().to_keystore();
        let key_type = "FSK";
        let version = "v1";

        let result = keystore.insert(key_type, version, KeyPair::mock());
        assert!(result.is_ok());

        let result = keystore.insert(key_type, version, KeyPair::mock());
        assert!(matches!(result, Err(KeyManagerError { .. })));
    }
}

#[cfg(test)]
impl KeyPair {
    fn mock() -> Self {
        Self {
            cached_private: OnceCell::new(),
            cached_public: OnceCell::new(),
            private_bytes: vec![],
            public_bytes: vec![],
            algorithm: "RS256".to_string(),
        }
    }
}
